import xarray as xr
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from tensorflow import keras

# ----------------------
from tensorflow.keras import layers
#我发现这个tensorflow.keras在tensorflow2.8被弃用了，最好用2.7版本。
#Python建议用3.9的，不建议用3.11和以上的。
#记得用conda配置清华源安装tensorflow以免出现依赖错误
# ----------------------

from tensorflow.keras.callbacks import EarlyStopping
from sklearn.metrics import mean_squared_error, r2_score
import matplotlib.pyplot as plt
import os

# ----------------------
# 1. 配置参数 (Configuration)
# ----------------------
wrfout_dir = 'path/to/your/wrfout/files'  # 使用wrfout文件
input_vars = [
    'T2',           # 2m 温度 (K)
    'PSFC',         # 地表气压 (Pa)
    'Q2',           # 2m 水汽混合比 (kg/kg)
    'U10',          # 10m U 风速 (m/s)
    'V10',          # 10m V 风速 (m/s)
    'RAINNC',       # 累积大尺度降水 (mm)
    'RAINSH',       # 累积对流降水 (mm)
    'SWDOWN',       # 地表向下短波辐射 (W/m^2)
    'LWDOWN',       # 地表向下长波辐射 (W/m^2)
    'O3',           # 对流层臭氧混合比 (ppmv)  (根据WRF-Chem 配置调整化学物种)
    'CO',           # 一氧化碳混合比 (ppmv)
    'NO',           # 一氧化氮混合比 (ppmv)
    'NO2',          # 二氧化氮混合比 (ppmv)
    'SO2'           # 二氧化硫混合比 (ppmv)
    # ... more something? maybe
]
output_vars = [
    'OA_VBS_1',     # 2DVBS 输出变量 1 (这里我们填第一个挥发性区间的有机气溶胶浓度)
    'OA_VBS_2',     # 2DVBS 输出变量 2 (这里我们填第二个挥发性区间的有机气溶胶浓度)
    # ... 
]

test_size = 0.2       # 测试集比例
validation_size = 0.2 # 验证集比例 (从训练集中划分)
random_state = 42     # 随机种子，保证结果可复现

# DNN 模型参数
dnn_layers = [128, 64, 32] # DNN 隐藏层神经元数量
activation_function = 'relu' # 激活函数
optimizer = 'adam'        # 优化器
loss_function = 'mse'      # 损失函数
metrics = ['mae', 'mse']   # 评估指标

# 训练参数
epochs = 100             # 训练轮数
batch_size = 32          # 批次大小
patience = 10            # EarlyStopping 的 patience 参数

# 保存模型路径
model_save_path = 'dnn_2dvbs_model.h5'

# ----------------------
# 2. 数据准备 (Data Preparation)
# ----------------------

def load_wrfout_data(wrfout_dir, input_vars, output_vars):
    """
    加载指定目录下的所有 wrfout 文件，提取输入和输出变量。
    """
    input_data_list = []
    output_data_list = []
    wrfout_files = [f for f in os.listdir(wrfout_dir) if f.startswith('wrfout_')] # 假设 wrfout 文件以 wrfout_ 开头

    for file in wrfout_files:
        file_path = os.path.join(wrfout_dir, file)
        try:
            ds = xr.open_dataset(file_path)

            # 检查变量是否存在
            for var in input_vars + output_vars:
                if var not in ds:
                    print(f"警告: 变量 '{var}' 在文件 '{file}' 中未找到，跳过该文件。")
                    ds.close() # 关闭数据集
                    continue # 跳到下一个文件

            input_data = ds[input_vars].to_dataframe().dropna() # 转换为 DataFrame 并去除 NaN 行
            output_data = ds[output_vars].to_dataframe().dropna()

            # 确保输入和输出数据有相同的索引 (时间步和网格点)
            common_index = input_data.index.intersection(output_data.index)
            input_data = input_data.loc[common_index]
            output_data = output_data.loc[common_index]

            input_data_list.append(input_data)
            output_data_list.append(output_data)
            ds.close() # 关闭数据集

        except Exception as e:
            print(f"读取文件 '{file}' 时出错: {e}")

    if not input_data_list:
        raise ValueError("没有成功加载任何 wrfout 文件并提取数据。请检查 wrfout 文件路径和变量名称。")

    # 合并所有文件的数据
    combined_input_data = pd.concat(input_data_list)
    combined_output_data = pd.concat(output_data_list)

    print(f"成功加载 {len(wrfout_files)} 个 wrfout 文件的数据。")
    return combined_input_data, combined_output_data


input_df, output_df = load_wrfout_data(wrfout_dir, input_vars, output_vars)

print("输入数据 (Input Data) 示例:")
print(input_df.head())
print("\n输出数据 (Output Data) 示例:")
print(output_df.head())


# 划分训练集、验证集和测试集
X_train_val, X_test, y_train_val, y_test = train_test_split(
    input_df, output_df, test_size=test_size, random_state=random_state
)
X_train, X_val, y_train, y_val = train_test_split(
    X_train_val, y_train_val, test_size=validation_size / (1 - test_size), random_state=random_state
) # 调整验证集比例


print(f"训练集大小: {X_train.shape[0]}")
print(f"验证集大小: {X_val.shape[0]}")
print(f"测试集大小: {X_test.shape[0]}")


# 数据标准化
input_scaler = StandardScaler()
output_scaler = StandardScaler() # 可以选择是否标准化输出变量，取决于咱们的需求

X_train_scaled = input_scaler.fit_transform(X_train) # 训练集 fit 和 transform
X_val_scaled = input_scaler.transform(X_val)       # 验证集 transform
X_test_scaled = input_scaler.transform(X_test)      # 测试集 transform

y_train_scaled = output_scaler.fit_transform(y_train) # 训练集 fit 和 transform
y_val_scaled = output_scaler.transform(y_val)       # 验证集 transform
y_test_scaled = output_scaler.transform(y_test)      # 测试集 transform


# ----------------------
# 3. 构建 DNN 模型 (Build DNN Model)
# ----------------------

model = keras.Sequential()
model.add(layers.Input(shape=(X_train_scaled.shape[1],))) # 输入层

# 添加隐藏层
for neurons in dnn_layers:
    model.add(layers.Dense(neurons, activation=activation_function))

model.add(layers.Dense(len(output_vars))) # 输出层，输出维度与 output_vars 数量一致 (线性激活)

# 打印模型结构
model.summary()

# 编译模型
model.compile(optimizer=optimizer, loss=loss_function, metrics=metrics)


# EarlyStopping 回调
early_stopping = EarlyStopping(
    monitor='val_loss',
    patience=patience,
    restore_best_weights=True,
    verbose=1
)

# ----------------------
# 4. 训练模型 (Train Model)
# ----------------------
history = model.fit(
    X_train_scaled, y_train_scaled,
    validation_data=(X_val_scaled, y_val_scaled),
    epochs=epochs,
    batch_size=batch_size,
    callbacks=[early_stopping],
    verbose=1
)


# ----------------------
# 5. 评估模型 (Evaluate Model)
# ----------------------

# 在测试集上评估
loss, mae, mse = model.evaluate(X_test_scaled, y_test_scaled, verbose=0)
rmse = np.sqrt(mse) # 计算 RMSE

print(f"测试集 Loss: {loss:.4f}")
print(f"测试集 MAE: {mae:.4f}")
print(f"测试集 MSE: {mse:.4f}")
print(f"测试集 RMSE: {rmse:.4f}")


# 预测测试集
y_pred_scaled = model.predict(X_test_scaled)

# 反标准化 (如果输出变量做了标准化)
y_pred = output_scaler.inverse_transform(y_pred_scaled)
y_true = output_scaler.inverse_transform(y_test_scaled)


# 计算 R-squared
r2 = r2_score(y_true, y_pred)
print(f"测试集 R-squared: {r2:.4f}")


# 可视化训练历史 (Loss 曲线)
plt.figure(figsize=(10, 6))
plt.plot(history.history['loss'], label='Training Loss')
plt.plot(history.history['val_loss'], label='Validation Loss')
plt.title('Training and Validation Loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
plt.show()


# 可视化预测结果 vs 真实值 (针对第一个输出变量)
output_var_name = output_vars[0] # 例如, 'OA_VBS_1'
plt.figure(figsize=(8, 8))
plt.scatter(y_true[:, 0], y_pred[:, 0], alpha=0.5) # 假设 y_true 和 y_pred 是 NumPy 数组
plt.xlabel(f'真实值 ({output_var_name})')
plt.ylabel(f'预测值 ({output_var_name})')
plt.title(f'真实值 vs 预测值 ({output_var_name})')
plt.plot([y_true[:, 0].min(), y_true[:, 0].max()], [y_true[:, 0].min(), y_true[:, 0].max()], 'k--', lw=2) # 对角线
plt.axis('equal')
plt.grid(True)
plt.show()


# ----------------------
# 6. 保存模型 (Save Model)
# ----------------------
model.save(model_save_path)
print(f"模型已保存到: {model_save_path}")


print("DNN 模型训练和评估完成！")